# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
# pylint: disable=unidiomatic-typecheck
"""Utility to lift subgraphs."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections

from tensorflow.python.framework import func_graph
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import resource_variable_ops


def _graph_inputs(op):
  return [x.op for x in op.inputs] + list(op.control_inputs)


def _as_operation(op_or_tensor):
  if isinstance(op_or_tensor, ops.Tensor):
    return op_or_tensor.op
  return op_or_tensor


class UnliftableError(Exception):
  """Raised if a Tensor cannot be lifted from the graph."""
  pass


def _constant_inputs(op_or_tensor):
  return all(_as_operation(i).type == u"Const"
             and not _as_operation(i).control_inputs
             for i in _graph_inputs(_as_operation(op_or_tensor)))


def _path_from(from_op, tensor, sources):
  """Find one path from `from_op` to `tensor`, ignoring `sources`.

  Args:
    from_op: A `tf.Operation`.
    tensor: A `tf.Operation` or `tf.Tensor`.
    sources: A list of `tf.Tensor`.

  Returns:
    A python string containing the path, or "??" if none is found.
  """
  visited_ops = set([x.op for x in sources])
  ops_to_visit = [_as_operation(tensor)]
  some_op_output = {}
  while ops_to_visit:
    op = ops_to_visit.pop()
    if op in visited_ops:
      continue
    visited_ops.add(op)
    if op == from_op:
      path_op = op
      path = [path_op]
      final_op = _as_operation(tensor)
      while path_op != final_op:
        path_op = some_op_output[path_op]
        path.append(path_op)
      return " <- ".join(["%s (%s)" % (x.name, x.type) for x in reversed(path)])
    else:
      for inp in _graph_inputs(op):
        if inp not in visited_ops and inp not in sources:
          some_op_output[inp] = op
          ops_to_visit.append(inp)
  return "??"


def _map_subgraph(init_tensor, sources, disallowed_placeholders, visited_ops,
                  op_outputs, add_sources):
  """Walk a Graph and capture the subgraph between init_tensor and sources.

  Note: This function mutates visited_ops and op_outputs.

  Arguments:
    init_tensor:  A Tensor or Operation where the subgraph terminates.
    sources:  A set of Tensors where subgraph extraction should stop.
    disallowed_placeholders: An optional set of ops which may not appear in the
      lifted graph. Defaults to all placeholders.
    visited_ops: A set of operations which were visited in a prior pass.
    op_outputs: A defaultdict containing the outputs of an op which are to be
      copied into the new subgraph.
    add_sources: A boolean indicating whether placeholders which are not in
      sources should be allowed.

  Returns:
    The set of placeholders upon which init_tensor depends and are not in
    sources.

  Raises:
    UnliftableError: if init_tensor depends on a placeholder which is not in
      sources and add_sources is False.
  """
  ops_to_visit = [_as_operation(init_tensor)]
  extra_sources = set()
  while ops_to_visit:
    op = ops_to_visit.pop()
    if op in visited_ops:
      continue
    visited_ops.add(op)

    should_raise = False
    if disallowed_placeholders is not None and op in disallowed_placeholders:
      should_raise = True
    elif op.type == "Placeholder":
      if disallowed_placeholders is None and not add_sources:
        should_raise = True
      extra_sources.update(op.outputs)

    if should_raise:
      raise UnliftableError(
          "Unable to lift tensor %s because it depends transitively on "
          "placeholder %s via at least one path, e.g.: %s"
          % (repr(init_tensor), repr(op), _path_from(op, init_tensor, sources)))
    for inp in _graph_inputs(op):
      op_outputs[inp].add(op)
      if inp not in visited_ops and inp not in (sources or extra_sources):
        ops_to_visit.append(inp)

  return extra_sources


def _copy_non_source(op, graph, op_map):
  """Copy an op directly to a given graph.

  This function assumes that all of the inputs to an op have already been
  copied.

  Args:
    op: The op to be copied.
    graph: The destination graph.
    op_map: A dict mapping ops and tensors in the old graph to the new one.
  """
  copied_inputs = [op_map[x] for x in op.inputs]
  copied_control_inputs = [op_map[x] for x in op.control_inputs]
  with ops.control_dependencies(copied_control_inputs), ops.device(op.device):
    copied_op = graph.create_op(
        op_type=op.type,
        inputs=copied_inputs,
        dtypes=[x.dtype for x in op.outputs],
        attrs=op.node_def.attr,
        name=op.name)
  op_map[op] = copied_op
  for i, o in enumerate(op.outputs):
    op_map[o] = copied_op.outputs[i]


def _copy_source(s, graph, op_map, handle_captures, inverse_captures):
  """Create a source in a graph based on a Tensor from a different graph.

  This function creates a placeholder analog of `s` in a graph with the
  following behavior:

  1) If s is a captured Tensor or Variable and handle_captures is set to True,
     simply capture it in the new graph as well.

  2) If s is a PlaceholderWithDefault whose default is a constant, preserve
     said default in the new graph.

  3) When applicable, copy resource variable metadata from `s` to the newly
     created placeholder.

  Args:
    s: The source of interest.
    graph: The destination graph.
    op_map: A dict mapping ops and tensors in the old graph to the new one.
    handle_captures: A boolean indicating whether to re-capture s in the new
      graph or simply create a vanilla placeholder.
    inverse_captures: A dict mapping s back to the Tensor or Variable that it
      captures.
  """
  if handle_captures and s in inverse_captures:
    copied_placeholder = graph.capture(inverse_captures[s], name=s.op.name)
  elif s.op.type == "PlaceholderWithDefault" and _constant_inputs(s):
    # Copy the default value to the graph.
    default_value = s.op.inputs[0]
    _copy_non_source(op=default_value.op, graph=graph, op_map=op_map)

    with ops.device(s.op.device):
      copied_placeholder = array_ops.placeholder_with_default(
          input=op_map[default_value], shape=s.shape, name=s.op.name)
  else:
    with ops.device(s.op.device):
      copied_placeholder = array_ops.placeholder(
          dtype=s.dtype, shape=s.shape, name=s.op.name)

  base_handle = resource_variable_ops.get_resource_handle_data(s)
  if base_handle.shape_and_type:
    resource_variable_ops._set_handle_shapes_and_types(  # pylint: disable=protected-access
        copied_placeholder,
        base_handle,
        graph_mode=True)

  op_map[s] = copied_placeholder


def lift_to_graph(init_tensors, graph, sources=None,
                  disallowed_placeholders=None, add_sources=False,
                  handle_captures=False, base_graph=None):
  """Copies the tensor and all its inputs recursively to the outer graph.

  Args:
    init_tensors: The Tensor to lift.
    graph: The graph to lift to.
    sources: Optional sequence of nodes to start from. If omitted the whole
      subgraph which feeds into `init_tensor` is lifted.
    disallowed_placeholders: An optional set of ops which may not appear in the
      lifted graph. Defaults to all placeholders.
    add_sources: A boolean indicating whether placeholders which are not in
      sources should be allowed.
    handle_captures: A boolean indicating whether to re-capture s in the new
      graph or simply create a vanilla placeholder.
    base_graph: The graph from which to lift ops. This will be inferred if not
      specified.

  Returns:
    A mapping from ops in the current default graph to ops in `graph`.

  Raises:
    UnliftableError: If a placeholder blocks lifting.
  """
  variable_init_tensors = {i for i in init_tensors if isinstance(
      i, resource_variable_ops.ResourceVariable)}
  init_tensors = set(init_tensors).difference(variable_init_tensors)
  base_graph = base_graph or list(init_tensors)[0].graph

  # Check that the initializer does not depend on any placeholders.
  sources = set(sources or [])
  visited_ops = set([x.op for x in sources])
  op_outputs = collections.defaultdict(set)

  # First we extract the subgraph between init_tensors and sources.
  for init_tensor in init_tensors:
    sources.update(_map_subgraph(
        init_tensor=init_tensor,
        sources=sources,
        disallowed_placeholders=disallowed_placeholders,
        visited_ops=visited_ops,
        op_outputs=op_outputs,
        add_sources=add_sources))

  # Topologically sort the nodes we've extracted. Now we know how many of their
  # outputs are part of this subgraph.
  ops_to_copy = []
  marked_ops = set([])
  ops_to_visit = [_as_operation(t) for t in init_tensors
                  if not op_outputs[_as_operation(t)]]
  while ops_to_visit:
    op = ops_to_visit.pop()
    if op in marked_ops:
      continue
    marked_ops.add(op)
    ops_to_copy.append(op)
    for inp in _graph_inputs(op):
      if (all(x in marked_ops for x in op_outputs[inp]) and
          inp not in sources):
        ops_to_visit.append(inp)

  # When lifting from one FuncGraph to another, we will need to capture the
  # relevant tensors as well.
  captures = collections.OrderedDict()
  if (isinstance(base_graph, func_graph.FuncGraph) and
      isinstance(graph, func_graph.FuncGraph)):
    captures = base_graph.captures
  inverse_captures = {v: k for k, v in captures.items()}

  # ops_to_copy now holds a reverse topologically sorted list of ops which
  # ends in the initializer. We copy those to the outermost graph and
  # build the initialization op there.
  with graph.as_default():
    op_map = {i: i for i in variable_init_tensors}  # Pass through variables.
    source_ops = set()
    for s in sources:
      source_ops.add(s.op)
      _copy_source(s=s, graph=graph, op_map=op_map,
                   handle_captures=handle_captures,
                   inverse_captures=inverse_captures)

    for op in reversed(ops_to_copy):
      if op in source_ops:
        continue

      _copy_non_source(op=op, graph=graph, op_map=op_map)

    return op_map
